#!/usr/bin/env python
# coding: utf-8

# | paddlepaddle | python|
# | -------- | -------- |
# | 2.4.0     | 3.7.4     |

# ### 1、环境配置

# In[1]:


# # 克隆PaddleSeg仓库, 执行一次就行
# from IPython import get_ipython
#
# get_ipython().system('git clone https://gitee.com/paddlepaddle/PaddleSeg')
#
#
# # In[2]:
#
#
# # 创建文件夹 执行一次就行
# get_ipython().system(' mkdir data')
# # 解压数据集 执行一次就行
# get_ipython().system(' unzip -oq ./data/data253280/安吉遥感语义分割数据集.zip -d data')


# ### 2、数据集可视化

# #### 2.1 可视化配置

# In[3]:


# 导库
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt

# 类别标签的颜色配置

color_map = {
    0: [0, 0, 0],        # '未知': 类别 0
    1: [215, 200, 185],  # '裸地': 类别 1
    2: [131, 194, 56],   # '草地': 类别 2
    3: [241, 165, 180],  # '构筑': 类别 3
    4: [210, 216, 201],  # '道路': 类别 4
    5: [49, 173, 105],   # '林地': 类别 5
    6: [163, 214, 245],  # '水域': 类别 6
    7: [248, 208, 114],  # '耕地': 类别 7
    8: [229, 103, 102]   # '房屋': 类别 8
}


def label2color(lable):
    '''
    将单通道灰度label, 根据颜色配置转为三通道彩色
    '''
    # 将图像转换为 numpy 数组
    label_array = np.array(lable)
    # 创建彩色图像
    color_label_array = np.zeros((label_array.shape[0], label_array.shape[1], 3), dtype=np.uint8)

    # 将单通道标签映射为彩色图像
    for label_value, color in color_map.items():
        color_label_array[label_array == label_value] = color
    return color_label_array


# #### 2.2 读取文件并可视化

# In[4]:


# 读取图像
image_path = '../data/images/AJ1_100.jpg'
label_path = '../data/labels/AJ1_100.png'

image = Image.open(image_path)
lable = Image.open(label_path)
print(f"image path: {image_path}, image size: {image.size}")
print(f"label path: {label_path}, label size: {lable.size}")

# 可视化
fig, axes = plt.subplots(1, 2, figsize=(1 * 5, 1 * 5))
fig.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0.01, wspace=0.01)
axes[0].imshow(image)
axes[0].axis("off")
axes[0].set_title("image")
axes[1].imshow(label2color(lable))
axes[1].axis("off")
axes[1].set_title("to color label")
plt.show()


# 注意数据集格式，详情可参考[自定义数据集](https://gitee.com/paddlepaddle/PaddleSeg/blob/release/2.8/docs/data/marker/marker_cn.md)。<br>
# 本文挂载的数据集格式如下，详情可见[遥感影像【语义分割】](https://aistudio.baidu.com/datasetdetail/253280)。
# ```
# data
# |─labels
# |   ├─AJ1_1.png
# |   ├─AJ1_2.png
# |─images
# |   ├─AJ1_1.jpg
# |   ├─AJ1_2.jpg
# |─train.txt
# |─val.txt
# |─test.txt
#
# ......
# train.txt如下：
# images/AJ2_1456.jpg labels/AJ2_1456.png
# images/AJ2_2188.jpg labels/AJ2_2188.png
#
# val.txt如下：
# images/AJ2_3046.jpg labels/AJ2_3046.png
# images/AJ2_1753.jpg labels/AJ2_1753.png
#
# test.txt如下：
# images/AJ2_2173.jpg labels/AJ2_2173.png
# ```

# ### 3、模型构建与训练

# In[5]:


# 导入需要用到的库
import sys
import os
import random
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from tqdm.notebook import tqdm


# In[6]:


# 定义全局变量

# 数据集存放目录
DATA_ROOT = '../data/'
# 训练集file_list文件路径
TRAIN_PATH= '../data/train.txt'
# 验证集file_list文件路径
VAL_PATH= '../data/val.txt'
# 测试集file_list文件路径
TEST_PATH= '../data/test.txt'
# 模型训练和visualdl日志文件的保存根路径
EXP_DIR = '../output'
# 测试集预测结果保存根路径
TEST_SAVE_DIR = '../output/results'
# 训练迭代次数（1 iter = 1 batch , 即进行一次前向传播和反向梯度下降）
# total iters = (total sample / batch size) * EPOCHS。若想使用EPOCHS，则需要知道total sample数量
ITERS = 1 # 这里只训练了不到2 epoch。样本数量6000左右，训练集5315，5315 // 4 = 1328（iter）。即1epoch可迭代1328次。
# EPOCHS = 100 # 暂时没使用到
# 单卡batch size
BACH_SIZE = 1
# 类别数(背景也算，即classes + background)
NUM_CLASSES = 9
# 图像大小
SIZE = (512, 512)
# 模型保存的间隔iter步数
SAVE_ITERS = 10
# 打印日志的间隔iter步数
LOG_ITERS = 1
# 用于异步读取数据的进程数量，大于等于1时开启子进程读取数据
NUM_WORKERS = 0


# In[7]:

# print(get_ipython()) # get_ipython().run_line_magic() 是 IPython 的魔法命令接口，只能在 IPython 交互式环境（如 Jupyter Notebook）中运行。在普通 Python 脚本或 IDE（如 PyCharm、VSCode）中调用时，get_ipython() 会返回 None，从而报错
# 跳转到PaddleSeg目录下，因为下面要使用克隆的源码中的paddleseg包
# get_ipython().run_line_magic('cd', 'PaddleSeg/')


# In[8]:


# 设置1张可用的卡（如果启动环境是GPU模式的情况下可运行，CPU环境不可运行）
# get_ipython().system(' export CUDA_VISIBLE_DEVICES=0')


# In[9]:


# 构建数据集
from paddleseg.datasets import Dataset
import paddleseg.transforms as T

# The transforms must be a list!
train_transforms = [
    T.Resize(target_size=SIZE),
    T.RandomHorizontalFlip(),
    T.Normalize()
]

val_transforms = [
    T.Resize(target_size=SIZE),
    T.Normalize()
]

train_dataset = Dataset(
    mode='train',
    dataset_root=DATA_ROOT,
    transforms=train_transforms,
    num_classes=NUM_CLASSES,
    train_path=TRAIN_PATH,
)

val_dataset = Dataset(
    mode='val',
    dataset_root=DATA_ROOT,
    transforms=val_transforms,
    num_classes=NUM_CLASSES,
    val_path=VAL_PATH,
)


# In[ ]:


# 构建模型
import paddle
from paddleseg.models import DeepLabV3
from paddleseg.models.backbones import ResNet18_vd, ResNet34_vd, ResNet50_vd, ResNet101_vd, ResNet152_vd
from paddleseg.models.losses import CrossEntropyLoss

# pretrained参数，使用paddle预训练的deeplabv3_resnet50模型
model = DeepLabV3(
    num_classes=NUM_CLASSES,
    backbone=ResNet50_vd(), # currently support Resnet50_vd/Resnet101_vd/Xception65.
    pretrained="https://bj.bcebos.com/paddleseg/dygraph/cityscapes/deeplabv3_resnet50_os8_cityscapes_1024x512_80k/model.pdparams"
)

# 设置学习率、优化器
base_lr = 0.01
lr = paddle.optimizer.lr.PolynomialDecay(base_lr, power=0.9, decay_steps=3000, end_lr=0)
optimizer = paddle.optimizer.SGD(lr, parameters=model.parameters(), weight_decay=4.0e-5)

# 设置损失函数
losses = {}
losses['types'] = [CrossEntropyLoss()] * 1
losses['coef'] = [1]* 1


# In[ ]:


# 开始训练
from paddleseg.core import train

# train(
#     model=model,
#     train_dataset=train_dataset,
#     val_dataset=val_dataset,
#     optimizer=optimizer,
#     losses=losses,
#     save_dir=EXP_DIR, # 模型和visualdl日志文件的保存根路径
#     iters=ITERS, # 训练迭代次数
#     batch_size=BACH_SIZE, # 单卡batch size
#     save_interval=SAVE_ITERS, # 模型保存的间隔步数
#     log_iters=LOG_ITERS, # 打印日志的间隔步数
#     num_workers=NUM_WORKERS, # 用于异步读取数据的进程数量， 大于等于1时开启子进程读取数据
#     use_vdl=True, # 是否开启visualdl记录训练数据
# )

# 开启visualdl后，可通过可视化功能，监测训练过程


# ### 4、模型评估与预测

# #### 4.1 模型评估

# In[ ]:


# 加载模型
import paddle
from paddleseg.models import DeepLabV3
from paddleseg.models.backbones import ResNet18_vd, ResNet34_vd, ResNet50_vd, ResNet101_vd, ResNet152_vd


model = DeepLabV3(
    num_classes=NUM_CLASSES,
    backbone=ResNet50_vd(), # currently support Resnet50_vd/Resnet101_vd/Xception65.
    pretrained=None
)

# 设置模型参数
model_path = '../output/best_model/model.pdparams'
if model_path:
    para_state_dict = paddle.load(model_path)
    model.set_dict(para_state_dict)
    print('Loaded trained params of model successfully')
else:
    raise ValueError('The model_path is wrong: {}'.format(model_path))


# In[ ]:


# 对测试集评估（也可评估验证集，这里只评估了测试集）

import paddleseg.transforms as T
from paddleseg.datasets import Dataset
from paddleseg.core import evaluate

# 构建测试集
# The transforms must be a list!
test_transforms = [
    T.Resize(target_size=SIZE),
    T.Normalize()
]

test_dataset = Dataset(
    mode='val', # 评估模式
    dataset_root=DATA_ROOT,
    transforms=test_transforms,
    num_classes=NUM_CLASSES,
    val_path=TEST_PATH, # 评估测试集
)

# 评估测试集
evaluate(model, test_dataset)


# ```
# 由上可知测试集评估结果如下：
# ```
#
# ###### 总体数据
#
# | Images | mIoU | Acc | Kappa | Dice |
# | -------- | -------- | -------- |-------- | -------- |
# | 665     | 0.5931     | 0.8185     | 0.7626     | 0.7305     |
#
#
# ###### 每一类数据
#
# | Class| 未知 | 裸地| 草地| 构筑物 | 道路| 林地| 水域| 耕地 | 房屋|
# | -------- | -------- | -------- | -------- | -------- | -------- | -------- | -------- | -------- | -------- |
# | Class label | 0     | 1     | 2     | 3     | 4     | 5     | 6     | 7     | 8     |
# | Class IoU     | 0.5812| 0.373| 0.5463| 0.2988|  0.6272|  0.8699|  0.7247|  0.6914|  0.6256|
# | Class Precision|  0.8196 | 0.4734 | 0.6953 | 0.5994 | 0.791  | 0.9095 | 0.8701 | 0.8347 | 0.7112|
# | Class Recall|  0.6665 | 0.6377|  0.7183 | 0.3734|  0.7517|  0.9523|  0.8127|  0.801 |  0.8387|
#
#

# #### 4.2 模型预测（用predict API方式）

# In[ ]:


import os


def get_image_list(data_root, test_path):
    """
    Get image list from the test path
    """

    image_list = []
    with open(test_path, 'r') as file:
        lines = file.readlines()
        for line in lines:
            image_path = os.path.join(data_root , line.split(' ')[0])
            if os.path.isfile(image_path):
                image_list.append(image_path)

    if len(image_list) == 0:
        raise RuntimeError('There are not image file in `--image_path`')

    return image_list

# 测试集图像文件目录、文件名列表
image_dir = os.path.join(DATA_ROOT, 'images')
image_list = get_image_list(DATA_ROOT, TEST_PATH)
# print(image_list)
# print(image_dir)


# In[ ]:


# 预测测试集
import paddleseg.transforms as T
from paddleseg.core import predict


# 用predict API测试时，一定要是T.Compose类型
test_transforms = T.Compose([
    T.Resize(target_size=SIZE),
    T.Normalize()
])

predict(
    model,
    model_path=model_path,  # 模型路径
    transforms=test_transforms, # transform.Compose, 对输入图像进行预处理
    image_list=image_list, # list, 待预测的图像路径列表。
    image_dir=image_dir, # str, 待预测的图片所在目录
    save_dir=TEST_SAVE_DIR #str, 结果输出路径
)


# ```
# 预测结果保存位置
# ./output/results/added_prediction
# ./output/results/pseudo_color_prediction
# ```
#
# | 原始图片 | 叠加预测结果 | 预测结果 |
# | -------- | -------- | -------- |
# | ![](https://ai-studio-static-online.cdn.bcebos.com/0d3fe9208eaa4845ad99dcdb8b49215f11f3e31c3bfb45e8b50c0fc55d71fbd1)     | ![](https://ai-studio-static-online.cdn.bcebos.com/4860dfab05d04f0cb1b033059d9f0ed760452877a90147efa48031827981f557)    | ![](https://ai-studio-static-online.cdn.bcebos.com/21554e7a223a43e3894b83efef2487c58cc974a979f144c2a0e8d1714916005a)     |
#
#
#
#
#

# #### 4.3 模型预测（非API方式，随机预测测试集并可视化结果）

# In[ ]:


# 可视化配置
import numpy as np

# 类别标签的颜色配置
color_map = {
    0: [0, 0, 0],        # '未知': 类别 0
    1: [215, 200, 185],  # '裸地': 类别 1
    2: [131, 194, 56],   # '草地': 类别 2
    3: [241, 165, 180],  # '构筑物':  类别 3
    4: [210, 216, 201],  # '道路': 类别 4
    5: [49, 173, 105],   # '林地': 类别 5
    6: [163, 214, 245],  # '水域': 类别 6
    7: [248, 208, 114],  # '耕地': 类别 7
    8: [229, 103, 102]   # '房屋': 类别 8
}

def label2color(lable):
    '''
    将单通道灰度label, 根据颜色配置转为三通道彩色
    '''
    # 将图像转换为 numpy 数组
    label_array = np.array(lable)
    # 创建彩色图像
    color_label_array = np.zeros((label_array.shape[0], label_array.shape[1], 3), dtype=np.uint8)

    # 将单通道标签映射为彩色图像
    for label_value, color in color_map.items():
        color_label_array[label_array == label_value] = color
    return color_label_array


# In[ ]:


# 构建测试集
import paddleseg.transforms as T
from paddleseg.datasets import Dataset

# 这个数据增强一定要是list类型！！！ 上面用predict API测试时是T.Compose类型
test_transforms = [
    T.Resize(target_size=SIZE),
    T.Normalize()
]

test_dataset = Dataset(
    mode='test',
    dataset_root=DATA_ROOT,
    transforms=test_transforms,
    num_classes=NUM_CLASSES,
    test_path=TEST_PATH,
)


# In[ ]:


# 可视化
import random
import paddle
from paddleseg.models import BiSeNetV2
from paddleseg.models import UNet
import paddleseg.transforms as T
from paddleseg.core import infer
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from paddleseg.core import infer

# get_ipython().run_line_magic('matplotlib', 'inline')

# 需要展示的样本个数
num_imgs_to_show = 4
# 随机抽取样本
chosen_indices = random.sample(range(len(test_dataset)), k=num_imgs_to_show)
# print(chosen_indices)

# 读取输入影像
image_paths = [test_dataset.file_list[idx][0] for idx in chosen_indices]
# 读取真值标签
label_paths = [test_dataset.file_list[idx][1] for idx in chosen_indices]

# 获取模型预测输出
preds = []
with paddle.no_grad():
    model.eval()
    for idx in chosen_indices:
        image = test_dataset[idx]['img']
        # C,H,W -> 1,C,H,W
        image = paddle.to_tensor(image).unsqueeze(0)
        # 预测推理
        pred = infer.inference(model, image)
        # 对pred进行argmax操作, 并转换为NumPy数组
        pred = paddle.argmax(pred, axis=1).numpy()
        # 调整数组形状为SIZE, 并将数据类型转换为np.uint8
        pred = np.reshape(pred, SIZE).astype(np.uint8)
        # 将预测结果转为彩色
        preds.append(label2color(pred))

# 绘制图片
fig, axes = plt.subplots(num_imgs_to_show, 3, figsize=(1 * 8, 1 * 8))
fig.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0.01, wspace=0.01)

# 给每列添加标题
axes[0, 0].set_title('Input Image')
axes[0, 1].set_title('Ground Truth Label')
axes[0, 2].set_title('Predicted Label')

for i in range(num_imgs_to_show):
    axes[i, 0].imshow(Image.open(image_paths[i]))
    axes[i, 0].axis('off')
    axes[i, 1].imshow(label2color(Image.open(label_paths[i])))
    axes[i, 1].axis('off')
    axes[i, 2].imshow(preds[i])
    axes[i, 2].axis('off')
plt.show()

